// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/vector/lazy_vector.h>
#include <pollux/common/base/runtime_metrics.h>
#include <pollux/common/memory/raw_vector.h>
#include <pollux/common/time/cpu_wall_timer.h>
#include <pollux/vector/complex_vector.h>
#include <pollux/vector/decoded_vector.h>
#include <pollux/vector/selectivity_vector.h>

namespace kumo::pollux {
    namespace {
        // Convenience class to record cpu and wall time from construction, updating
        // thread local stats at destruction, including input bytes of the vector passed
        // as parameter.
        class LazyIoStatsRecorder {
        public:
            LazyIoStatsRecorder(VectorPtr *vector) : vector_(vector) {
            }

            ~LazyIoStatsRecorder() {
                auto cpuDelta = timer_.elapsed();
                addThreadLocalRuntimeStat(
                    LazyVector::kWallNanos,
                    RuntimeCounter(cpuDelta.wallNanos, RuntimeCounter::Unit::kNanos));
                addThreadLocalRuntimeStat(
                    LazyVector::kCpuNanos,
                    RuntimeCounter(cpuDelta.cpuNanos, RuntimeCounter::Unit::kNanos));

                if (*vector_) {
                    addThreadLocalRuntimeStat(
                        LazyVector::kInputBytes,
                        RuntimeCounter(
                            (*vector_)->estimate_flat_size(), RuntimeCounter::Unit::kBytes));
                }
            }

        private:
            DeltaCpuWallTimeStopWatch timer_;
            VectorPtr *vector_;
        };
    } // namespace

    void VectorLoader::load(
        RowSet rows,
        ValueHook *hook,
        vector_size_t resultSize,
        VectorPtr *result) { {
            LazyIoStatsRecorder recorder(result);
            loadInternal(rows, hook, resultSize, result);
        }
        if (hook) {
            // Record number of rows loaded directly into ValueHook bypassing
            // materialization into vector. This counter can be used to understand
            // whether aggregation pushdown is happening or not.
            addThreadLocalRuntimeStat("loadedToValueHook", RuntimeCounter(rows.size()));
        }
    }

    void VectorLoader::load(
        const SelectivityVector &rows,
        ValueHook *hook,
        vector_size_t resultSize,
        VectorPtr *result) {
        if (rows.isAllSelected()) {
            const auto &indices = DecodedVector::consecutiveIndices();
            POLLUX_DCHECK(!indices.empty());
            if (rows.end() <= indices.size()) {
                load(
                    RowSet(&indices[rows.begin()], rows.end() - rows.begin()),
                    hook,
                    resultSize,
                    result);
                return;
            }
        }
        raw_vector<vector_size_t> positions(rows.countSelected());
        simd::indicesOfSetBits(
            rows.allBits(), rows.begin(), rows.end(), positions.data());
        load(positions, hook, resultSize, result);
    }

    VectorPtr LazyVector::slice(vector_size_t offset, vector_size_t length) const {
        POLLUX_CHECK(isLoaded(), "Cannot take slice on unloaded lazy vector");
        POLLUX_DCHECK(vector_);
        return vector_->slice(offset, length);
    }

    //   static
    // This gurantee that vectors are loaded but it does not gurantee that vector
    // will be free of any nested lazy vector i.e, references to Loaded lazy vectors
    // are not replaced with their underlying loaded vector.
    void LazyVector::ensureLoadedRows(
        const VectorPtr &vector,
        const SelectivityVector &rows) {
        if (is_lazy_not_loaded(*vector)) {
            DecodedVector decoded;
            SelectivityVector baseRows;
            ensureLoadedRows(vector, rows, decoded, baseRows);
        } else {
            // Even if LazyVectors have been loaded, via some other path, this
            // is needed for initializing wrappers.
            vector->loaded_vector();
        }
    }

    namespace {
        // Given a SelectivityVector 'rows', updates 'baseRows' selecting the rows
        // in the base vector that should be loaded.
        void selectBaseRowsToLoad(
            const DecodedVector &decoded,
            SelectivityVector &baseRows,
            const SelectivityVector &rows) {
            POLLUX_DCHECK(
                decoded.base()->encoding() == VectorEncoding::Simple::ROW ||
                decoded.base()->encoding() == VectorEncoding::Simple::LAZY);

            auto deselectNullsIdentity = [&]() {
                if (decoded.base()->raw_nulls()) {
                    baseRows.deselectNulls(
                        decoded.base()->raw_nulls(), rows.begin(), rows.end());
                }
            };

            if (decoded.isIdentityMapping() && rows.isAllSelected()) {
                baseRows.resizeFill(rows.end(), true);
                deselectNullsIdentity();
                return;
            }

            if (decoded.isIdentityMapping()) {
                baseRows.resizeFill(rows.end(), false);
                baseRows.select(rows);
                deselectNullsIdentity();
            } else if (decoded.isConstantMapping()) {
                baseRows.resizeFill(decoded.index(0) + 1, false);
                baseRows.setValid(decoded.index(0), true);
            } else {
                baseRows.resizeFill(decoded.base()->size(), false);
                rows.applyToSelected([&](vector_size_t row) {
                    if (!decoded.is_null_at(row)) {
                        baseRows.setValid(decoded.index(row), true);
                    }
                });
            }

            baseRows.updateBounds();
        }
    } // namespace

    // static
    void LazyVector::ensureLoadedRowsImpl(
        const VectorPtr &vector,
        DecodedVector &decoded,
        const SelectivityVector &rows,
        SelectivityVector &baseRows) {
        if (decoded.base()->encoding() != VectorEncoding::Simple::LAZY) {
            if (decoded.base()->encoding() == VectorEncoding::Simple::ROW &&
                is_lazy_not_loaded(*decoded.base())) {
                auto *row_vector = decoded.base()->as_unchecked<RowVector>();
                selectBaseRowsToLoad(decoded, baseRows, rows);
                DecodedVector decodedChild;
                SelectivityVector childRows;
                for (auto child: row_vector->children()) {
                    decodedChild.decode(*child, baseRows, false);
                    ensureLoadedRowsImpl(child, decodedChild, baseRows, childRows);
                }
                row_vector->updateContainsLazyNotLoaded();
                vector->loaded_vector();
            }
            return;
        }
        // Base vector is lazy.
        auto baseLazyVector = decoded.base()->as_unchecked<LazyVector>();

        if (!baseLazyVector->isLoaded()) {
            // Create rowSet.
            raw_vector<vector_size_t> rowNumbers;
            RowSet rowSet;
            if (decoded.isConstantMapping()) {
                rowNumbers.push_back(decoded.index(rows.begin()));
                rowSet = RowSet(rowNumbers);
            } else if (decoded.isIdentityMapping()) {
                if (rows.isAllSelected()) {
                    auto iota = pollux::iota(rows.end(), rowNumbers);
                    rowSet = RowSet(iota, rows.end());
                } else {
                    rowNumbers.resize(rows.end());
                    rowNumbers.resize(simd::indicesOfSetBits(
                        rows.asRange().bits(), 0, rows.end(), rowNumbers.data()));
                    rowSet = RowSet(rowNumbers);
                }
            } else {
                selectBaseRowsToLoad(decoded, baseRows, rows);
                rowNumbers.resize(baseRows.end());
                rowNumbers.resize(simd::indicesOfSetBits(
                    baseRows.asRange().bits(), 0, baseRows.end(), rowNumbers.data()));

                rowSet = RowSet(rowNumbers);
            }

            baseLazyVector->load(rowSet, nullptr);
        }

        // The loaded vector can itself also be lazy, so we load recursively.
        if (is_lazy_not_loaded(*baseLazyVector->vector_)) {
            // We do not neeed to decode all rows.
            selectBaseRowsToLoad(decoded, baseRows, rows);
            decoded.decode(*baseLazyVector->vector_, baseRows, false);
            SelectivityVector nestedRows;
            ensureLoadedRowsImpl(
                baseLazyVector->vector_, decoded, baseRows, nestedRows);
        }

        // An explicit call to loaded_vector() is necessary to allow for proper
        // initialization of dictionaries, sequences, etc. on top of lazy vector
        // after it has been loaded, even if loaded via some other path.
        vector->loaded_vector();
    }

    // static
    void LazyVector::ensureLoadedRows(
        const VectorPtr &vector,
        const SelectivityVector &rows,
        DecodedVector &decoded,
        SelectivityVector &baseRows) {
        decoded.decode(*vector, rows, false);
        ensureLoadedRowsImpl(vector, decoded, rows, baseRows);
    }

    void LazyVector::validate(const VectorValidateOptions &options) const {
        if (isLoaded() || options.loadLazy) {
            auto loaded_vector = this->loaded_vector();
            POLLUX_CHECK_NOT_NULL(loaded_vector);
            loaded_vector->validate(options);
        }
    }

    void LazyVector::load(RowSet rows, ValueHook *hook) const {
        POLLUX_CHECK(!allLoaded_, "A LazyVector can be loaded at most once");

        allLoaded_ = true;
        if (rows.empty()) {
            vector_ = BaseVector::create_null_constant(type_, size(), pool_);
            return;
        }

        // The loader expect structs to be pre-allocated.
        if (!vector_ && type_->kind() == TypeKind::ROW) {
            vector_ = BaseVector::create(type_, rows.back() + 1, pool_);
        }

        // We do not call prepare_for_reuse on vector_ in purpose, the loader handles
        // that.
        loader_->load(rows, hook, size(), &vector_);
        if (!hook) {
            POLLUX_CHECK_GE(vector_->size(), size());
        }
    }

    void LazyVector::loadVectorInternal() const {
        if (!allLoaded_) {
            if (!vector_) {
                vector_ = BaseVector::create(type_, 0, pool_);
            }
            SelectivityVector allRows(BaseVector::length_);
            loader_->load(allRows, nullptr, size(), &vector_);
            POLLUX_CHECK_NOT_NULL(vector_);
            if (vector_->encoding() == VectorEncoding::Simple::LAZY) {
                vector_ = vector_->as_unchecked<LazyVector>()->loaded_vector_shared();
            } else {
                // If the load produced a wrapper, load the wrapped vector.
                vector_->loaded_vector();
            }
            allLoaded_ = true;
            const_cast<LazyVector *>(this)->BaseVector::nulls_ = vector_->nulls_;
            if (BaseVector::nulls_) {
                const_cast<LazyVector *>(this)->BaseVector::rawNulls_ =
                        BaseVector::nulls_->as<uint64_t>();
            }
        } else {
            POLLUX_CHECK_NOT_NULL(vector_);
        }
    }
} // namespace kumo::pollux
